"""
Nonlinear PDE solver using locELM method.

Solves nonlinear PDEs using NLSQ-perturb and Newton-LLSQ methods.
"""

import numpy as np
from scipy.optimize import least_squares
from scipy.linalg import lstsq
from typing import Callable, Dict, Optional, Tuple
from .linear import LinearPDESolver


class NonlinearPDESolver(LinearPDESolver):
    """Solver for nonlinear PDEs using locELM."""

    def solve_nonlinear(self,
                       residual_function: Callable,
                       initial_guess: Optional[np.ndarray] = None,
                       method: str = 'nlsq-perturb',
                       nlsq_options: Optional[Dict] = None) -> np.ndarray:
        """
        Solve nonlinear PDE using specified method.

        Parameters
        ----------
        residual_function : Callable
            Function that computes residual vector given weights
        initial_guess : Optional[np.ndarray]
            Initial guess for weights
        method : str
            Solution method ('nlsq-perturb' or 'newton-llsq')
        nlsq_options : Optional[Dict]
            Options for nonlinear solver

        Returns
        -------
        np.ndarray
            Solution weights
        """
        if method == 'nlsq-perturb':
            return self._solve_nlsq_perturb(residual_function, initial_guess, nlsq_options)
        elif method == 'newton-llsq':
            return self._solve_newton_llsq(residual_function, initial_guess, nlsq_options)
        else:
            raise ValueError(f"Unknown method: {method}")

    def _solve_nlsq_perturb(self,
                           residual_function: Callable,
                           initial_guess: Optional[np.ndarray] = None,
                           options: Optional[Dict] = None) -> np.ndarray:
        """
        Solve using NLSQ-perturb method (Algorithm 1 in paper).

        Parameters
        ----------
        residual_function : Callable
            Function that computes residual
        initial_guess : Optional[np.ndarray]
            Initial guess
        options : Optional[Dict]
            Solver options

        Returns
        -------
        np.ndarray
            Solution weights
        """
        # Default options
        default_options = {
            'delta': 0.5,
            'cost_threshold': 1e-3,
            'max_sub_iterations': 20,
            'ftol': 1e-12,
            'xtol': 1e-12
        }
        if options:
            default_options.update(options)

        delta = default_options['delta']
        cost_threshold = default_options['cost_threshold']
        max_sub_iter = default_options['max_sub_iterations']

        # Initial guess
        if initial_guess is None:
            x0 = np.zeros(self.networks.n_total_params)
        else:
            x0 = initial_guess.copy()

        # Initial solve
        result = least_squares(
            residual_function,
            x0,
            ftol=default_options['ftol'],
            xtol=default_options['xtol']
        )

        x_best = result.x.copy()
        cost_best = result.cost

        # Check if converged
        if cost_best < cost_threshold:
            self.networks.set_all_weights(x_best)
            return x_best

        # Sub-iteration with perturbations
        for i in range(max_sub_iter):
            # Generate random perturbation
            xi_1 = np.random.uniform(0, 1)
            delta_1 = xi_1 * delta

            # Random perturbation vector
            Delta_x = np.random.uniform(-delta_1, delta_1, x_best.shape)

            # Generate random mixing parameter
            xi_2 = np.random.uniform(0, 1)
            y0 = xi_2 * x_best + Delta_x

            # Solve with perturbed initial guess
            result = least_squares(
                residual_function,
                y0,
                ftol=default_options['ftol'],
                xtol=default_options['xtol']
            )

            # Update if better
            if result.cost < cost_best:
                x_best = result.x.copy()
                cost_best = result.cost

            # Check convergence
            if cost_best < cost_threshold:
                break

        # Set final weights
        self.networks.set_all_weights(x_best)
        return x_best

    def _solve_newton_llsq(self,
                          residual_function: Callable,
                          initial_guess: Optional[np.ndarray] = None,
                          options: Optional[Dict] = None) -> np.ndarray:
        """
        Solve using Newton-LLSQ method (Newton iterations + linear least squares).

        Parameters
        ----------
        residual_function : Callable
            Function that computes residual
        initial_guess : Optional[np.ndarray]
            Initial guess
        options : Optional[Dict]
            Solver options

        Returns
        -------
        np.ndarray
            Solution weights
        """
        # Default options
        default_options = {
            'max_iter': 100,
            'tol': 1e-6,
            'finite_diff_eps': 1e-7
        }
        if options:
            default_options.update(options)

        max_iter = default_options['max_iter']
        tol = default_options['tol']
        eps = default_options['finite_diff_eps']

        # Initial guess
        if initial_guess is None:
            W = np.zeros(self.networks.n_total_params)
        else:
            W = initial_guess.copy()

        # Newton iterations
        for k in range(max_iter):
            # Compute residual
            G = residual_function(W)

            # Check convergence
            norm_G = np.linalg.norm(G)
            if norm_G < tol:
                break

            # Compute Jacobian using finite differences
            m = len(G)
            n = len(W)
            J = np.zeros((m, n))

            for j in range(n):
                W_perturbed = W.copy()
                W_perturbed[j] += eps
                G_perturbed = residual_function(W_perturbed)
                J[:, j] = (G_perturbed - G) / eps

            # Solve linear least squares: J * Delta_W = -G
            Delta_W, residual, rank, s = lstsq(J, -G, lapack_driver='gelsy')

            # Update
            W = W + Delta_W

        # Set final weights
        self.networks.set_all_weights(W)
        return W

    def build_nonlinear_residual(self,
                                nonlinear_operator: Callable,
                                source_term: Callable,
                                boundary_conditions: Dict,
                                output_dim: int = 1) -> Callable:
        """
        Build residual function for nonlinear PDE.

        Parameters
        ----------
        nonlinear_operator : Callable
            Nonlinear operator function F(u, u_x, u_y, ...)
        source_term : Callable
            Source term f(x)
        boundary_conditions : Dict
            Boundary conditions
        output_dim : int
            Number of output components

        Returns
        -------
        Callable
            Residual function that takes weights and returns residual vector
        """
        def residual_func(weights: np.ndarray) -> np.ndarray:
            # Set network weights
            self.networks.set_all_weights(weights)

            # Compute PDE residual
            pde_residual = []
            for subdomain_idx in range(self.domain.n_total):
                points = self.collocation_points[subdomain_idx]
                network = self.networks.get_network(subdomain_idx)

                # Get network output and derivatives
                u = network(points)

                # Compute nonlinear operator
                F_u = nonlinear_operator(network, points, subdomain_idx)

                # Compute linear part (if any, assume included in operator)
                # Source term
                f = source_term(points)
                if f.ndim == 1:
                    f = f.reshape(-1, 1)

                # Residual: F(u, u_x, u_y) - f = 0
                for comp in range(output_dim):
                    for i in range(points.shape[0]):
                        res = F_u[i, comp] - f[i, min(comp, f.shape[1]-1)]
                        pde_residual.append(res)

            # Compute boundary residual
            bc_residual = []
            for subdomain_idx in range(self.domain.n_total):
                boundaries = self.continuity.interface_manager.get_boundary_interfaces(subdomain_idx)
                network = self.networks.get_network(subdomain_idx)

                for boundary in boundaries:
                    if boundary['dimension'] in boundary_conditions:
                        bc_spec = boundary_conditions[boundary['dimension']]

                        if boundary['type'] in bc_spec:
                            bc_func = bc_spec[boundary['type']]
                            points = self.collocation_points[subdomain_idx]

                            # Filter boundary points
                            dim = boundary['dimension']
                            coord = boundary['coordinate']
                            mask = np.abs(points[:, dim] - coord) < 1e-10
                            boundary_points = points[mask]

                            if boundary_points.shape[0] > 0:
                                u = network(boundary_points)
                                g = bc_func(boundary_points)
                                if g.ndim == 1:
                                    g = g.reshape(-1, 1)

                                for comp in range(output_dim):
                                    for i in range(boundary_points.shape[0]):
                                        res = u[i, comp] - g[i, min(comp, g.shape[1]-1)]
                                        bc_residual.append(res)

            # Compute continuity residual
            A_cont, b_cont = self.continuity.build_continuity_equations(
                self.collocation_points, output_dim
            )
            cont_residual = A_cont @ weights - b_cont

            # Combine all residuals
            residual = np.hstack([pde_residual, bc_residual, cont_residual])
            return residual

        return residual_func
